{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data as data\n",
    "\n",
    "import random\n",
    "import numbers\n",
    "from PIL import Image, ImageMath\n",
    "import os\n",
    "import os.path\n",
    "import numpy as np\n",
    "import struct\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "import faiss\n",
    "import time\n",
    "\n",
    "import pcl\n",
    "from PIL import Image, ImageDraw\n",
    "import faiss\n",
    "\n",
    "import timeit\n",
    "from pyquaternion import Quaternion\n",
    "\n",
    "%matplotlib qt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KNNBuilder_GPU:\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "        self.dimension = 3\n",
    "        \n",
    "        # we need only a StandardGpuResources per GPU\n",
    "        self.res = faiss.StandardGpuResources()\n",
    "        self.res.setTempMemoryFraction(0.1)\n",
    "        self.flat_config = faiss.GpuIndexFlatConfig()\n",
    "        self.flat_config.device = 0\n",
    "        \n",
    "    def build_nn_index(self, database):\n",
    "        '''\n",
    "        :param database: numpy array of Nx3\n",
    "        :return: Faiss index, in CPU\n",
    "        '''\n",
    "        index = faiss.GpuIndexFlatL2(self.res, self.dimension, self.flat_config)  # dimension is 3\n",
    "        index.add(database)\n",
    "        return index\n",
    "    \n",
    "    def search_nn(self, index, query, k):\n",
    "        '''\n",
    "        :param index: Faiss index\n",
    "        :param query: numpy array of Nx3\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        D, I = index.search(query, k)\n",
    "        return D, I\n",
    "    \n",
    "    def self_build_search(self, x):\n",
    "        '''\n",
    "\n",
    "        :param x: numpy array of Nxd\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        x = np.ascontiguousarray(x, dtype=np.float32)\n",
    "        index = self.build_nn_index(x)\n",
    "        D, I = self.search_nn(index, x, self.k)\n",
    "        return D, I\n",
    "    \n",
    "\n",
    "class KNNBuilder:\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "        self.dimension = 3\n",
    "\n",
    "    def build_nn_index(self, database):\n",
    "        '''\n",
    "        :param database: numpy array of Nx3\n",
    "        :return: Faiss index, in CPU\n",
    "        '''\n",
    "        index = faiss.IndexFlatL2(self.dimension)  # dimension is 3\n",
    "        index.add(database)\n",
    "        return index\n",
    "\n",
    "    def search_nn(self, index, query, k):\n",
    "        '''\n",
    "        :param index: Faiss index\n",
    "        :param query: numpy array of Nx3\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        D, I = index.search(query, k)\n",
    "        return D, I\n",
    "\n",
    "    def self_build_search(self, x):\n",
    "        '''\n",
    "\n",
    "        :param x: numpy array of Nxd\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        x = np.ascontiguousarray(x, dtype=np.float32)\n",
    "        index = self.build_nn_index(x)\n",
    "        D, I = self.search_nn(index, x, self.k)\n",
    "        return D, I\n",
    "    \n",
    "\n",
    "class PCSampler:\n",
    "    def __init__(self, leaf_size, minimum_pc_num):       \n",
    "        self.leaf_size = leaf_size\n",
    "        self.minimum_pc_num = minimum_pc_num\n",
    "    \n",
    "    def sample_pc(self, pc, leaf_size):\n",
    "        '''\n",
    "        :param pc: input numpy array of Nx3\n",
    "        :return: sampled_pc of Mx3\n",
    "        '''\n",
    "        cloud = pcl.PointCloud(pc)\n",
    "        sor = cloud.make_voxel_grid_filter()\n",
    "        sor.set_leaf_size(leaf_size, leaf_size, leaf_size)\n",
    "        cloud_filtered = sor.filter()\n",
    "        sampled_pc = np.asarray(cloud_filtered)\n",
    "        \n",
    "        return sampled_pc\n",
    "    \n",
    "    def sample_pc_wrapper(self, pc):\n",
    "        '''\n",
    "        ensure that the sampled pc is more than a certain amount\n",
    "        '''\n",
    "        retry_counter = 0\n",
    "        \n",
    "        sampled_pc = self.sample_pc(pc, self.leaf_size)\n",
    "        while sampled_pc.shape[0] < self.minimum_pc_num:\n",
    "            retry_counter += 1\n",
    "            leaf_size = self.leaf_size - 0.04*retry_counter\n",
    "            if leaf_size <= 0:\n",
    "                break\n",
    "            sampled_pc = self.sample_pc(pc, leaf_size)\n",
    "        \n",
    "        return sampled_pc\n",
    "    \n",
    "    \n",
    "def axisEqual3D(ax):\n",
    "    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])\n",
    "    sz = extents[:,1] - extents[:,0]\n",
    "    centers = np.mean(extents, axis=1)\n",
    "    maxsize = max(abs(sz))\n",
    "    r = maxsize/2\n",
    "    for ctr, dim in zip(centers, 'xyz'):\n",
    "        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_bin(file_path, num_cols=6):\n",
    "    pc_np = np.fromfile(file_path, dtype=np.float32)\n",
    "    pc_np = np.reshape(pc_np, (-1, num_cols))\n",
    "    \n",
    "    # remove the surface normals, I will compute it myself using python-pcl\n",
    "    return pc_np[:, 0:3]\n",
    "\n",
    "\n",
    "def load_test_gt_txt(txt_path):\n",
    "    f = open(txt_path, 'r')\n",
    "    lines_list = f.readlines()\n",
    "    \n",
    "    sample_num = 828  # first line is header\n",
    "    pos_idx_list = []  # a list that stores the nearby sample indexes for each sample\n",
    "    for i in range(sample_num):\n",
    "        pos_idx_list.append([])\n",
    "    \n",
    "    dataset = []\n",
    "    for i, line_str in enumerate(lines_list):\n",
    "        if i == 0:\n",
    "            # skip the headline\n",
    "            continue\n",
    "        # convert each line to a dict\n",
    "        line_splitted_list = line_str.split()\n",
    "        try:\n",
    "            assert len(line_splitted_list) == 11\n",
    "        except Exception:\n",
    "            print('Invalid line.')\n",
    "            print(i)\n",
    "            print(line_splitted_list)\n",
    "            continue\n",
    "        \n",
    "        anc_idx = int(line_splitted_list[0].strip())\n",
    "        pos_idx = int(line_splitted_list[1].strip())\n",
    "        t = [float(line_splitted_list[4].strip()), \n",
    "             float(line_splitted_list[5].strip()), \n",
    "             float(line_splitted_list[6].strip())]  # tx, ty, tz\n",
    "        q = [float(line_splitted_list[7].strip()), \n",
    "             float(line_splitted_list[8].strip()), \n",
    "             float(line_splitted_list[9].strip()),\n",
    "             float(line_splitted_list[10].strip())]  # quaternion: w, x, y, z\n",
    "        \n",
    "#         print(anc_idx)\n",
    "#         print(pos_idx)\n",
    "#         print(t)\n",
    "#         print(q)\n",
    "\n",
    "        data = {'anc_idx': anc_idx, 'pos_idx': pos_idx, 't': t, 'q': q}\n",
    "        dataset.append(data)\n",
    "        \n",
    "        \n",
    "        # build negative sample\n",
    "        pos_idx_list[anc_idx].append(pos_idx)\n",
    "        pos_idx_list[pos_idx].append(anc_idx)\n",
    "        \n",
    "        \n",
    "    f.close()\n",
    "    \n",
    "#     print(pos_idx_list)\n",
    "    \n",
    "    \n",
    "    # build negative sample\n",
    "    # set random seed to ensure the same result\n",
    "    random.seed(1)\n",
    "    for i, data in enumerate(dataset):\n",
    "        anc_idx = data['anc_idx']\n",
    "        \n",
    "        while True:\n",
    "            rand_idx = random.randint(0, sample_num-1)  # a<=x<=b\n",
    "            if not (rand_idx in pos_idx_list[anc_idx]):\n",
    "                data['neg_idx'] = rand_idx\n",
    "                break\n",
    "    \n",
    "    # reset reandom seed\n",
    "    random.seed()\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_test_gt_txt('/ssd/dataset/oxford/test_models_20k/groundtruths.txt')\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 - min_pc_num 19062\n",
      "100 - min_pc_num 18905\n",
      "200 - min_pc_num 18905\n",
      "300 - min_pc_num 18905\n",
      "400 - min_pc_num 18905\n",
      "500 - min_pc_num 18905\n",
      "600 - min_pc_num 18905\n",
      "700 - min_pc_num 18905\n",
      "800 - min_pc_num 18905\n",
      "minimum point number: 18905\n"
     ]
    }
   ],
   "source": [
    "# get surface normal\n",
    "def Surface_normals(cloud):\n",
    "    ne = cloud.make_NormalEstimation()\n",
    "    tree = cloud.make_kdtree()\n",
    "    ne.set_SearchMethod(tree)\n",
    "#     ne.set_RadiusSearch(2)\n",
    "    ne.set_KSearch(9)\n",
    "    cloud_normals = ne.compute()\n",
    "    return cloud_normals\n",
    "\n",
    "# 1. voxel grid filer\n",
    "# 2. apply outlier filter\n",
    "# 3. calculate surface normal and curvature\n",
    "def process_pc(pc_np):\n",
    "#     print('origin')\n",
    "#     print(pc_np.shape)\n",
    "    \n",
    "    # 1. voxel grid filer\n",
    "    sampler = PCSampler(leaf_size=0.2, minimum_pc_num=20480)\n",
    "    pc_np = sampler.sample_pc_wrapper(pc_np)\n",
    "#     print('after voxel grid filter')\n",
    "#     print(pc_np.shape)\n",
    "    \n",
    "    # 2. apply outlier filter\n",
    "    fil = pcl.PointCloud(pc_np).make_statistical_outlier_filter()\n",
    "    fil.set_mean_k(50)\n",
    "    fil.set_std_dev_mul_thresh(2)\n",
    "    pc_np = np.asarray(fil.filter())\n",
    "#     print('after outlier filter')\n",
    "#     print(pc_np.shape)\n",
    "    \n",
    "#     fig = plt.figure()\n",
    "#     ax = Axes3D(fig)\n",
    "#     ax.scatter(pc_np[:,0].tolist(), pc_np[:,1].tolist(), pc_np[:,2].tolist(), s=0.1, c=[0.5,0.5,0.5])\n",
    "#     axisEqual3D(ax)\n",
    "\n",
    "#     plt.ion()\n",
    "#     plt.show()\n",
    "    \n",
    "    # 3. calculate surface normal and curvature\n",
    "    assert pc_np.shape[1] == 3\n",
    "    cloud = pcl.PointCloud(pc_np)    \n",
    "    sn = Surface_normals(cloud)\n",
    "    sn_np = np.asarray(sn.to_array(), dtype=np.float32)  # Nx4, nx,ny,nz,curvature\n",
    "    \n",
    "    output_np = np.concatenate((pc_np, sn_np), axis=1)  # Nx7\n",
    "    \n",
    "    return output_np\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# convert bin to npy\n",
    "bin_folder = '/ssd/dataset/oxford/test_models_20k'\n",
    "npy_folder = '/ssd/dataset/oxford/test_models_20k_np'\n",
    "\n",
    "min_pc_num = 1000000\n",
    "content = os.listdir(bin_folder)\n",
    "for i, filename in enumerate(content):\n",
    "    if not ('.bin' in filename):\n",
    "        continue\n",
    "    \n",
    "    # load bin\n",
    "    pc_np = load_bin(os.path.join(bin_folder, filename))\n",
    "    \n",
    "    # process pc\n",
    "    pc_np = process_pc(pc_np)\n",
    "\n",
    "    # save pc as numpy array\n",
    "    npy_file_name = filename[0:-3] + 'npy'\n",
    "    np.save(os.path.join(npy_folder, npy_file_name), pc_np)\n",
    "\n",
    "    if pc_np.shape[0] < min_pc_num:\n",
    "        min_pc_num = pc_np.shape[0]\n",
    "        \n",
    "    if i % 100 == 0:\n",
    "        print('%d - min_pc_num %d' % (i, min_pc_num))\n",
    "\n",
    "print('minimum point number: %d' % min_pc_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test registration\n",
    "folder = '/ssd/dataset/oxford/test_models_20k_np'\n",
    "dataset = load_test_gt_txt('/ssd/dataset/oxford/test_models_20k_np/groundtruths.txt')\n",
    "print('length of dataset: %d' % len(dataset))\n",
    "for i, data in enumerate(dataset):\n",
    "    if i<80:\n",
    "        continue\n",
    "    \n",
    "    anc_idx = data['anc_idx']\n",
    "    pos_idx = data['pos_idx']\n",
    "    neg_idx = data['neg_idx']\n",
    "    t = np.asarray(data['t'], dtype=np.float32)\n",
    "    q_list = data['q']\n",
    "    q = Quaternion(np.asarray(q_list))\n",
    "    \n",
    "    T = q.transformation_matrix\n",
    "    \n",
    "    T[0, 3] = t[0]\n",
    "    T[1, 3] = t[1]\n",
    "    T[2, 3] = t[2]\n",
    "    \n",
    "    T = T[0:3, :]  # 3x4\n",
    "    \n",
    "    anc_pc_np = np.load(os.path.join(folder, '%d.npy'%anc_idx))[:, 0:3]\n",
    "    pos_pc_np = np.load(os.path.join(folder, '%d.npy'%pos_idx))[:, 0:3]\n",
    "#     anc_pc_np = load_bin(os.path.join('/ssd/dataset/oxford/test_models_16384', '%d.bin'%anc_idx))[:, 0:3]\n",
    "#     pos_pc_np = load_bin(os.path.join('/ssd/dataset/oxford/test_models_16384', '%d.bin'%pos_idx))[:, 0:3]\n",
    "    \n",
    "    N = pos_pc_np.shape[0]\n",
    "    pos_pc_np_pad = np.concatenate((pos_pc_np, np.ones((N, 1))), axis=1)\n",
    "    print(pos_pc_np_pad.shape)\n",
    "    \n",
    "    tmp_pc_np = np.transpose(np.dot(T, np.transpose(pos_pc_np_pad)))\n",
    "    print(tmp_pc_np.shape)\n",
    "    \n",
    "    fig = plt.figure()\n",
    "    ax = Axes3D(fig)\n",
    "    ax.scatter(anc_pc_np[:,0].tolist(), anc_pc_np[:,1].tolist(), anc_pc_np[:,2].tolist(), s=0.1, c=[0.5,0.5,0.5])\n",
    "    ax.scatter(tmp_pc_np[:,0].tolist(), tmp_pc_np[:,1].tolist(), tmp_pc_np[:,2].tolist(), s=0.1, c=[1, 0, 0])\n",
    "#     ax.scatter(pos_pc_np[:,0].tolist(), pos_pc_np[:,1].tolist(), pos_pc_np[:,2].tolist(), s=0.1, c=[0, 0, 1])\n",
    "    axisEqual3D(ax)\n",
    "\n",
    "    plt.ion()\n",
    "    plt.show()\n",
    "    \n",
    "    print('anc: %d, pos: %d, neg: %d' % (anc_idx, pos_idx, neg_idx))\n",
    "    print(t)\n",
    "    print(q)\n",
    "    print(T)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save txt to pickle\n",
    "import pickle\n",
    "\n",
    "def save_obj(obj, name):\n",
    "    with open(name, 'wb') as f:\n",
    "        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "def load_obj(name):\n",
    "    with open(name, 'rb') as f:\n",
    "        return pickle.load(f)\n",
    "    \n",
    "save_obj(dataset, '/ssd/dataset/oxford/test_models_20k_np/groundtruths.pkl')\n",
    "dataset_read = load_obj('/ssd/dataset/oxford/test_models_20k_np/groundtruths.pkl')\n",
    "print(dataset_read)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert txt to that used in descriptor training\n",
    "import pickle\n",
    "\n",
    "def save_obj(obj, name):\n",
    "    with open(name, 'wb') as f:\n",
    "        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "def load_obj(name):\n",
    "    with open(name, 'rb') as f:\n",
    "        return pickle.load(f)\n",
    "    \n",
    "dataset = load_obj('/ssd/dataset/oxford/test_models_20k_np/groundtruths.pkl')\n",
    "dataset_descriptor = []\n",
    "dataset_descriptor_idx_set = set([])\n",
    "\n",
    "for i, item in enumerate(dataset):   \n",
    "    # process anc_idx\n",
    "    anc_idx = item['anc_idx']\n",
    "    if anc_idx not in dataset_descriptor_idx_set:\n",
    "        nonneg_set = set([])\n",
    "        nonneg_set.add(anc_idx)\n",
    "        for j, item_j in enumerate(dataset):\n",
    "            if item_j['anc_idx'] == anc_idx:\n",
    "                nonneg_set.add(item_j['pos_idx'])\n",
    "            if item_j['pos_idx'] == anc_idx:\n",
    "                nonneg_set.add(item_j['anc_idx'])\n",
    "        nonneg_idx_list = list(nonneg_set)\n",
    "        dataset_descriptor.append({'anc_idx': anc_idx, 'pos_idx_list': nonneg_idx_list})\n",
    "        dataset_descriptor_idx_set.add(anc_idx)\n",
    "    \n",
    "    # process pos_idx\n",
    "    anc_idx = item['pos_idx']\n",
    "    if anc_idx not in dataset_descriptor_idx_set:\n",
    "        nonneg_set = set([])\n",
    "        nonneg_set.add(anc_idx)\n",
    "        for j, item_j in enumerate(dataset):\n",
    "            if item_j['anc_idx'] == anc_idx:\n",
    "                nonneg_set.add(item_j['pos_idx'])\n",
    "            if item_j['pos_idx'] == anc_idx:\n",
    "                nonneg_set.add(item_j['anc_idx'])\n",
    "        nonneg_idx_list = list(nonneg_set)\n",
    "        dataset_descriptor.append({'anc_idx': anc_idx, 'pos_idx_list': nonneg_idx_list})\n",
    "        dataset_descriptor_idx_set.add(anc_idx)\n",
    "\n",
    "\n",
    "    \n",
    "# print(dataset_descriptor)\n",
    "print(len(dataset_descriptor))\n",
    "print(len(dataset_descriptor_idx_set))\n",
    "\n",
    "# for i in range(828):\n",
    "#     if i not in dataset_descriptor_idx_set:\n",
    "#         print(i)\n",
    "            \n",
    "save_obj(dataset_descriptor, '/ssd/dataset/oxford/test_models_20k_np/gt_descriptor_testing.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset_descriptor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
